import os
import xml.etree.ElementTree as ET
import glob
import tqdm


def cal_iou(obj_1, obj_2, iou_thr):
    lable_1 = obj_1[0]
    bbox_1 = obj_1[1:]
    #
    is_ok = False
    for i in obj_2:
        lable_i = i[0]
        bbox_i = i[1:]
        if lable_i != lable_1:
            continue
        area_1 = (bbox_1[2] - bbox_1[0]) * (bbox_1[3] - bbox_1[1])
        area_2 = (bbox_i[2] - bbox_i[0]) * (bbox_i[3] - bbox_i[1])
        #
        lt_x = max(bbox_1[0], bbox_i[0])
        lt_y = max(bbox_1[1], bbox_i[1])
        rb_x = min(bbox_1[2], bbox_i[2])
        rb_y = min(bbox_1[3], bbox_i[3])
        #
        overlap = (rb_x - lt_x) * (rb_y - lt_y)
        #
        union = area_1 + area_2 - overlap
        #
        iou = overlap / union
        #
        if iou > iou_thr:
            is_ok = True
            break
    
    return is_ok


if __name__ == '__main__':
    #
    path_1 = 'temp/temp1'   #文件夹1
    path_2 = 'temp/temp2'   #文件夹2
    #
    # iou域值
    iou_thr = 0.6
    #
    list_dir_1 = glob.glob(os.path.join(path_1, '*.xml'))   #os.listdir(path_1)
    list_dir_2 = os.listdir(path_2)

    # 不相似文件保存下来
    write_file = open('train.txt', 'w')

    # 开始遍历
    for filename in tqdm.tqdm(list_dir_1):
        filename = os.path.basename(filename)
        if filename.endswith('.xml'):
            #如果 文件夹1中对应xml文件 在 文件夹2 中没有，打印出来
            if filename not in list_dir_2:
                write_file.write(filename + '\n')
                print(filename, '\t中存在不相似')
                continue
            #
            #path_1中xml标注框信息
            root_1 = ET.parse(os.path.join(path_1, filename)).getroot()
            objects_1 = root_1.findall('object')
            #
            #path_2中xml标注框信息
            root_2 = ET.parse(os.path.join(path_2, filename)).getroot()
            objects_2 = root_2.findall('object')
            #
            #先判断两个文件的标注框个数，若个数不同，直接判为不相似
            if len(objects_1) != len(objects_2):
                write_file.write(filename)
                print(filename, '\t中存在不相似')
                continue
            #
            #遍历objects_1中的标注框 与 objects_2中的所有标注框 比较， 若类别相等，计算iou，否则跳过
            #
            #objects_2 --> annos_2
            annos_2 = []
            for obj_j in objects_2:
                lable = obj_j.find('name').text
                sub = obj_j.find('bndbox')  # 找到框的标注值并进行读取
                x1 = float(sub[0].text)
                y1 = float(sub[1].text)
                x2 = float(sub[2].text)
                y2 = float(sub[3].text)
                dict_j = [lable, x1, y1, x2, y2]
                annos_2.append(dict_j)
            #
            #objects_1
            for obj_i in objects_1:
                lable = obj_i.find('name').text
                sub = obj_i.find('bndbox')  # 找到框的标注值并进行读取
                x1 = float(sub[0].text)
                y1 = float(sub[1].text)
                x2 = float(sub[2].text)
                y2 = float(sub[3].text)
                dict_i = [lable, x1, y1, x2, y2]
                is_ok = cal_iou(dict_i, annos_2, iou_thr)
                #若is_ok为True，则obj_i 在 annos_2 有对应的标注框；否则存在不相似情况，并打印结果
                if not is_ok:
                    write_file.write(filename + '\n')
                    print(filename, '\t中存在不相似')
                    break
